import random
import time
from robomaster import robot
from robomaster import conn
from multi_robomaster import multi_robot
from robomaster import led, blaster
from collections import deque

import pos_nokov as pos
import logging
import sympy as sp
import math
import copy
import msvcrt
from interval import Interval
from position import *
from utils import *
import center_control as cc


# field_error=False
# field_range=[[-1.33,3.04],[-1.50,1.70]]
aim_x=Interval(-1.15, 3.00)
aim_y=Interval(-1.30,1.50)
safe_x=Interval(-1.32,2.76)
safe_y=Interval(-1.42,1.61)

rl=1.2 # 右路在y轴的数值
ml=-0.2 # 中路在y轴的数值
goal=pose(-1.1,ml,180)

adjust=0.8

passing=0
catching=0

def led_red_solid(robot_group):
    """红灯常亮"""
    robot_group.led.set_led(led.COMP_ALL, 255, 1, 1, led.EFFECT_ON)

def led_grey_solid(robot_group):
    """grey常亮"""
    robot_group.led.set_led(led.COMP_ALL, 75, 255, 255, led.EFFECT_ON)
    
def calc_catching_pose(b, t, aim=False, catch_phase1=False):
    '''Calc pose to catch ball. b: ball pose, t: target pose'''
    x,y=sp.symbols('x, y')
    distance=0.316 if catch_phase1 else 0.42 # distance from ball to pose
    # logging.debug(f'dist:{distance}')
    solutions=sp.nonlinsolve([(x-b.x)*(b.y-t.y)-(y-b.y)*(b.x-t.x),
                              (x-b.x)**2+(y-b.y)**2-distance**2],[x,y])
    for s in solutions:
        if aim: # the catching pose aim at t
            if not (s[0] in Interval(b.x, t.x) and s[1] in Interval(b.y,t.y)):
                return pose(s[0],s[1],math.atan2(t.y-b.y, t.x-b.x)*180/math.pi)
        else: # t is agent pose and catch directly
            # logging.debug(f'sx, {s[0]},sy: {s[1]}; bx: {b.x}, by: {b.y}; tx: {t.x}, ty: {t.y}')
            if (s[0]-b.x)*(t.x-b.x)>0 and (s[1]-b.y)*(t.y-b.y)>0:
                # logging.debug(f'solved pt: x:{s[0]},y:{s[1]}')
                return pose(s[0],s[1],math.atan2(b.y-t.y, b.x-t.x)*180/math.pi)
    
def catch_mobile(p, b, v):
    '''
    给定球的位置b、速度v、机器人位置p，决定是否夹爪
    return: bool
    '''
    # logging.debug(f'vx:{v.x},vy:{v.y}')
    tpred=get_len(b.x-p.x,b.y-p.y)/get_len(v.x,v.y)
    lerror=abs((b.x-p.x)*v.y-(b.y-p.y)*v.x)
    logging.debug(f'judge, time to collide: {tpred}, line error: {lerror}') # 预计碰撞时间
    # logging.debug(lerror) # 共线则接近0
    if lerror<0.05 and tpred<1.2:
        return True
    
class EP:
    def __init__(self, multi_robots, robot_group, id, role='stand_by'):
        self.role=role
        self.multi_robots=multi_robots
        self.chase_phase=0
        self.catch_phase=0
        self.pass_phase=0
        self.close_to_oppo=False
        self.robot_group=robot_group
        self.id=id
        self.ce=0 # closest enemy id
        self.aim_to_pass=False
        self.trans_pt=pose(0.5,0.5,5)
        if self.role=='stand_by':
            cc.sb.append(self.id)
        
    def init_arm(self, rg):
        rg.robotic_arm.moveto(0, -70).wait_for_completed()
        rg.gripper.open(100)
        time.sleep(2)
    
    def init_pos(self,rg):
        while True:
            p=pos.p
            v=pos.v
            v_best=get_rvo_vel_old(p.get_arr(), v.get_arr(),cc.pt.get_arr())
            self.trans_pt=trans_relative_co(pose(0,0,p[self.id].phi), pose(v_best[self.id].x,v_best[self.id].y))
            phi_pt=limit_pi(cc.pt[self.id].phi-p[self.id].phi)
            if dis_pose(cc.pt[self.id], p[self.id])<0.1 and phi_pt<3:
                # logging.debug(f'ep{self.id} current:{p.arr[self.id]},target:{cc.pt[self.id]}')
                break
            self.trans_pt.update(self.trans_pt.x, self.trans_pt.y, phi_pt)
            rg.chassis.drive_speed(math.tanh(self.trans_pt.x), -math.tanh(self.trans_pt.y/2), -45*(math.tanh(self.trans_pt.phi/45)), timeout=0.5)
            time.sleep(0.05)
        
        
    def stop(self, rg):
        rg.chassis.drive_wheels(w1=0, w2=0, w3=0, w4=0,timeout=1)
        time.sleep(0.01)
    
    def chase(self, rg):
        p=pos.p
        # 如果self.ready_to_catch为真，则以包住球为target pose
        # if self.aim_to_pass: # chase while aiming the stand_by teammate
        #     logging.info('chase aiming')
        #     cc.pt[self.id]=calc_catching_pose(p[0], p[cc.sb[0]], True, self.chase_phase==1)
        # else: # chase directly
        # logging.debug(f'ep{self.id} chase directly')
        cc.pt[self.id]=calc_catching_pose(p[0], p[self.id], False, self.chase_phase==1)
        # if self.chase_phase==0:
        #     v_best=get_rvo_vel(p.arr, v.get_arr(),cc.pt.arr)
        #     trans_pt=trans_relative_co(pose(0,0,p[self.id].phi), pose(v_best[self.id][0],v_best[self.id][1])) # 获得机器人坐标系下的目标坐标pt
        # elif self.chase_phase==1:
        self.trans_pt=trans_relative_co(p[self.id], cc.pt[self.id]) # 获得机器人坐标系下的目标坐标pt
        
        phi_pt=limit_pi(cc.pt[self.id].phi-p[self.id].phi)
        self.trans_pt.update(self.trans_pt.x, self.trans_pt.y, phi_pt)
        # logging.debug(f'self.trans_pt is {self.trans_pt}')
        rg.chassis.drive_speed(math.tanh(self.trans_pt.x/1.5), -math.tanh(self.trans_pt.y), -30*(math.tanh(self.trans_pt.phi/45)), timeout=1)
        time.sleep(0.1)
        
    def catch_still(self, rg):
        self.multi_robots.run([self.robot_group, self.stop])
        p=pos.p
        rg.gripper.close(100)
        time.sleep(0.5)
        if check_catched(p[self.id],p[0]):
            logging.info(f'ep{self.id} catched')
            self.chase_phase=3
            cc.bp=self.id
            # t1=time.time()
            # rg.robotic_arm.moveto(x=0, y=-60).wait_for_completed()
            # t2=time.time()
            # logging.debug(f'arm lift time:{t2-t1: .3f} s')
        else:
            logging.info(f'ep{self.id} catch failed')
            self.chase_phase=1
            rg.gripper.open(100)
            time.sleep(0.31) 
    
    def prepare_to_pass(self, rg):
        global passing
        p=pos.p
        passing=1
        self.trans_pt=get_passing_pose(p[self.id],p[cc.catcher],p[cc.eid])
        rg.chassis.drive_speed(self.trans_pt.x, self.trans_pt.y, -45*(math.tanh(self.trans_pt.phi/45)), timeout=1)
        time.sleep(0.1)
    
    def pass_ball(self,rg):
        global passing
        passing=1
        p=pos.p
        # t2=time.time()
        dphi=limit_pi(p[self.id].phi-math.atan2(p[cc.catcher].y-p[self.id].y,
                        p[cc.catcher].x-p[self.id].x)*180/math.pi)
        # logging.debug(f'dphi to pass:{dphi}')
        if abs(dphi)>2:
            rg.chassis.drive_speed(0,0,45*(math.tanh(dphi/40)))
            time.sleep(0.1)
        elif catching!=0:
            logging.info(f'ep{self.id} start passing2')
            rg.chassis.drive_speed(0,0,0)
            time.sleep(0.01)
            
            rg.chassis.drive_speed(2,0,0,timeout=2) # shoot
            time.sleep(0.2)
            # rg.robotic_arm.moveto(0,5).wait_for_completed() # move up arm
            
            rg.gripper.open(power=100)
            rg.chassis.drive_speed(2,0,0,timeout=2) # shoot
            time.sleep(0.1)
            rg.chassis.drive_speed(-2,0,0,timeout=2) # shoot
            time.sleep(0.2)
            rg.chassis.drive_speed(0,0,0) # stop
            time.sleep(0.01)
            passing=2
            self.chase_phase=0
            if cc.bp==self.id:# 防止提前被接住
                cc.bp=0
            passing=0
            
            # rg.chassis.move(x=-1.4,xy_speed=3).wait_for_completed()
            # rg.robotic_arm.moveto(0, -70).wait_for_completed()
            # t3=time.time()
            # logging.debug(f'ep{self.id} finished passing in {t3-t2}s')
            self.role='stand_by'
            cc.sb.append(self.id)
            logging.info(f'ep{self.id} finished passing, in stand_by, current sb: {cc.sb}')
        else:
            time.sleep(0.1)  
    
    def prepare_to_catch(self,rg):
        global bp
        p=pos.p
        dphi=angle_to_aim(p[self.id], p[cc.bp])-p[self.id].phi
        # logging.debug(f'ep{self.id} prepare to catch: rotate {dphi}')
        rg.chassis.drive_speed(0,0,-45*(math.tanh(limit_pi(dphi)/45)), timeout=1)
        time.sleep(0.1)
        
    def catch(self, rg):
        global adjust
        p=pos.p
        v=pos.v
        self.trans_pt=calc_closest_pose_on_v(p[self.id],p[0],v[0])
        rg.chassis.drive_speed(adjust*math.tanh(self.trans_pt.x/5), -adjust*math.tanh(self.trans_pt.y), -45*(math.tanh(self.trans_pt.phi/45)), timeout=1)
        if v[0].get_norm()>0.5:
            logging.debug(f'catching moving ball of speed {v[0].get_norm()}')
            if catch_mobile(p[self.id],p[0],v[0]):
                rg.gripper.close(power=100)
                time.sleep(0.3)
                if check_catched(p[self.id],p[0]):
                    logging.info('catched')
                    self.chase_phase=3
                    cc.bp=self.id
                    rg.robotic_arm.moveto(x=0, y=-60).wait_for_completed()
                else:
                    logging.info('catch failed')
                    self.chase_phase=1
                    rg.gripper.open(100)
                    time.sleep(0.25) 
        time.sleep(0.1)
    
    def attack(self, rg):
        global goal
        p=pos.p
        if self.id==3:
            cc.pt[self.id]=pose(-1.4, rl,180)
        elif self.id==1:
            cc.pt[self.id]=goal
        
        self.trans_pt=trans_relative_co(p[self.id], cc.pt[self.id])
        phi_pt=limit_pi(cc.pt[self.id].phi-p[self.id].phi)
        self.trans_pt.update(self.trans_pt.x, self.trans_pt.y, phi_pt)
        rg.chassis.drive_speed(adjust*math.tanh(self.trans_pt.x/5), -adjust*math.tanh(self.trans_pt.y), -45*(math.tanh(self.trans_pt.phi/45)), timeout=1)
        time.sleep(0.1)
    
    def forward(self, rg):
        p=pos.p
        v=pos.v
        
        if self.id==1:
            cc.pt[self.id]=pose(p[3].x, ml, 180)
        elif self.id==3:
            cc.pt[self.id]=pose(p[1].x, rl, 180)#angle_to_aim(p[self.id], p[cc.bp])
        # self.ce, sd=closest_enemy(p[self.id],p)
        self.ce, d=obstacle_enemy(p[self.id], cc.pt[self.id], p)
        if self.ce!=0:
            logging.info(f'obstacle detected for ep{self.id}: ep{self.ce}, {d}m to the route')
        # logging.debug(f'')
        if self.ce!=0 and abs(cc.pt[self.id].x-p[self.ce].x)<0.6:
            logging.info(f'ep{self.id} forwarding, met ep{self.ce} in {abs(cc.pt[self.id].x-p[self.ce].x)}')
            cc.pt[self.id]=pose(p[self.ce].x-0.7,cc.pt[self.id].y,180)
            v_best=get_rvo_vel_old(p.get_arr(), v.get_arr(),cc.pt.get_arr())
            self.trans_pt=trans_relative_co(pose(0,0,p[self.id].phi), pose(v_best[self.id].x,v_best[self.id].y))
            phi_pt=limit_pi(cc.pt[self.id].phi-p[self.id].phi)
            self.trans_pt.update(self.trans_pt.x, self.trans_pt.y, phi_pt)
        else:
            self.trans_pt=trans_relative_co(p[self.id], cc.pt[self.id]) # 获得机器人坐标系下的目标坐标pt
            phi_pt=limit_pi(cc.pt[self.id].phi-p[self.id].phi)
            self.trans_pt.update(self.trans_pt.x, self.trans_pt.y, phi_pt)
        rg.chassis.drive_speed(math.tanh(self.trans_pt.x/1.5), -math.tanh(self.trans_pt.y), -45*(math.tanh(self.trans_pt.phi/45)), timeout=1)
        time.sleep(0.1)
        
    def run(self,event):
        logging.basicConfig(format='%(asctime)s.%(msecs)03d-[line:%(lineno)d]'
                            '-%(levelname)s: %(message)s',
                            level=logging.DEBUG, datefmt='%H:%M:%S')
        global passing,catching
        cc.ready[self.id]=0
        self.multi_robots.run([self.robot_group, self.init_arm])
        event.wait()
        p=pos.p
        v=pos.v
        cc.pt=copy.deepcopy(p)
        # initial pose
        cc.pt[1]=pose(2.6,ml, 180)
        cc.pt[3]=pose(2.6,rl,180)
        self.multi_robots.run([self.robot_group, self.init_pos])
        logging.info(f'ep{self.id} reached init pos')
        self.multi_robots.run([self.robot_group, self.stop])
        cc.ready[self.id]=1
        while 0 in cc.ready:
            time.sleep(0.1)
        time.sleep(0.5)
        slow_c=0
        while not pos.done:
            p=pos.p
            try:
                # if pt.x in safe_x and pt.y in safe_y: 
                #     self.aim_to_pass=True
                # else:
                self.aim_to_pass=False
                if self.role=='chase':
                    if self.chase_phase==0:
                        # if -0.15<self.trans_pt.x<0.3 and abs(self.trans_pt.y)<0.06 and abs(self.trans_pt.phi)<5:
                        if check_inline(p[self.id], p[0]):    
                            logging.info(f'ep{self.id} chase phase 1')
                            self.chase_phase=1
                            continue
                        self.multi_robots.run([self.robot_group, self.chase])
                    elif self.chase_phase==1:
                        self.multi_robots.run([self.robot_group, self.chase])
                        if abs(self.trans_pt.x)<0.02 and abs(self.trans_pt.y)<0.02 and abs(self.trans_pt.phi)<5:
                            logging.info(f'ep{self.id} chase phase 2')
                            self.chase_phase=2
                            self.multi_robots.run([self.robot_group, self.stop])
                            continue    
                        
                    elif self.chase_phase==2:
                        self.multi_robots.run([self.robot_group, self.catch_still])
                    elif self.chase_phase==3:
                        logging.info(f'ep{self.id} attack')
                        self.role='attack'
                elif self.role=='attack':
                    cc.eid, sd=closest_enemy(p[self.id],p)
                    if sd<0.7:
                        if p[self.id].x-p[cc.eid].x>0.1:
                            self.multi_robots.run([self.robot_group, self.stop])
                            # self.ce=cc.eid
                            if not check_contain_only(cc.sb, self.id):
                                logging.info(f'ep{self.id} met ep{cc.eid} in {sd}m and pass_phase 0')
                                self.role='pass'
                                self.pass_phase=0
                                cc.catcher=cc.sb[0]
                    elif dis_pose(p[self.id], goal)<0.1:
                        logging.info(f'ep{self.id} reached goal')
                        break
                    else:
                        self.multi_robots.run([self.robot_group, self.attack])
                elif self.role=='pass':
                    if self.pass_phase==0:
                        # 先避开对手再传球
                        self.multi_robots.run([self.robot_group, self.prepare_to_pass])
                        if self.trans_pt.y==0:
                            logging.info(f'ep{self.id} pass phase 1')
                            self.pass_phase=1
                    elif self.pass_phase==1:
                        self.multi_robots.run([self.robot_group, self.pass_ball])
                
                elif self.role=='stand_by':
                    if passing!=0 and self.id==cc.sb[0] and dis_pose(cc.pt[self.id], p[self.id])<0.2:
                        logging.info(f'ep{self.id} try to catch')    
                        self.multi_robots.run([self.robot_group, self.stop])
                        catching=1
                        self.role='catch'
                        cc.sb.popleft()
                        
                        # logging.info(f'ep{self.id} catching dynamic vb:{v[0]}')
                        # logging.info(f'current stand_by list:{cc.sb}')
                    else:
                        self.multi_robots.run([self.robot_group, self.forward])
                elif self.role=='catch':
                    if passing==1:
                        self.multi_robots.run([self.robot_group, self.prepare_to_catch])
                    else:
                        # if np.linalg.norm(np.array([v[0].x,v[0].y]))>0.5:
                        #     slow_c=0
                        # else:
                        #     logging.debug(f'sc={slow_c}, ball slow with speed: {v[0]}')
                        #     slow_c+=1
                        #     # time.sleep(0.05)
                        # if slow_c<2:
                        #     self.multi_robots.run([self.robot_group, self.catch])
                        # else:
                        # logging.info(f'ball seems too slow, ep{self.id} chasing, phase 0')
                        self.role='chase'
                        self.chase_phase=0
                        catching=0
            except Exception as e:
                logging.error(f'Exception: {e}')
                break
            
            if (msvcrt.kbhit() and msvcrt.getch() == b'p'):
                print("ep %d Paused..." % self.id)
                while msvcrt.getch() != b'r':
                    time.sleep(0.5)
                    msvcrt.getch()
                print("ep %d resumed..." % self.id)
            if msvcrt.kbhit() and msvcrt.getch() == b't':
                break    
        pos.done=True
        print('over...')
        self.multi_robots.close()
